Skip to content

Persistent AG + GEMM kernels w/ HipKittens & UserBuffers - #707

Open
alextmagro wants to merge 4 commits into
devfrom
hipkittens_persistent
Open

Persistent AG + GEMM kernels w/ HipKittens & UserBuffers#707
alextmagro wants to merge 4 commits into
devfrom
hipkittens_persistent

Conversation

@alextmagro

Copy link
Copy Markdown
Contributor

Introduces persistent AG + GEMM kernels for the TN and NN layout -- 1.247x and 1.304x speedup on average across 24 training sizes.

Win is entirely in overlap -- GEMM speeds are 0.885x and 0.923x compared to hipBLASLt, so optimization opportunity available there. FSDP overlap is considered robust, but only tested with an intranode simulation.

Comment thread transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp Outdated
Comment thread transformer_engine/common/gemm/kittens/kittens_common.cpp Outdated
Comment thread transformer_engine/pytorch/module/base.py Outdated
Comment thread transformer_engine/pytorch/module/base.py Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/gemm.cpp Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_nn.cuh Outdated
Comment thread tests/pytorch/distributed/test_comm_gemm_overlap.py Outdated
Comment thread tests/pytorch/distributed/test_comm_gemm_overlap.py Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh
@ROCm ROCm deleted a comment from github-actions Bot Aug 18, 2026
@wenchenvincent wenchenvincent added the ci-level 3 CI test level 3 label Aug 19, 2026
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh
Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm.cpp Outdated
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh Outdated
Comment thread transformer_engine/pytorch/module/base.py
Comment thread tests/pytorch/distributed/test_comm_gemm_overlap.py
Comment thread transformer_engine/common/gemm/kittens/cdna4/fused_ag_gemm_tn.cuh Outdated
@alextmagro
alextmagro requested a review from wangye805 August 21, 2026 14:45
@aris134
aris134 self-requested a review August 21, 2026 18:24
Comment thread transformer_engine/pytorch/module/base.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci-level 3 CI test level 3

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants